msg_tool\scripts\artemis\ast/
dump.rs

1use super::types::*;
2use crate::utils::escape::*;
3use std::io::Write;
4
5struct LenChecker {
6    target_len: usize,
7    current_len: usize,
8}
9
10impl LenChecker {
11    fn new(target_len: usize) -> Self {
12        LenChecker {
13            target_len,
14            current_len: 0,
15        }
16    }
17
18    fn check(&mut self, value: &Value) -> bool {
19        match value {
20            Value::Float(f) => {
21                if f.fract() == 0.0 {
22                    self.current_len += format!("{:.1}", f).len();
23                } else {
24                    self.current_len += format!("{}", f).len();
25                }
26            }
27            Value::Int(i) => self.current_len += format!("{}", i).len(),
28            Value::Str(s) => {
29                self.current_len += s.len()
30                    + if lua_str_contains_need_escape(s) {
31                        4
32                    } else {
33                        2
34                    }
35            }
36            Value::KeyVal((k, v)) => {
37                if let Some(key) = k.as_str() {
38                    self.current_len += key.as_bytes().len()
39                        + if lua_key_contains_need_escape(key) {
40                            7
41                        } else {
42                            3
43                        };
44                } else {
45                    self.current_len += 1; // for '['
46                    if !self.check(k) {
47                        return false;
48                    }
49                    self.current_len += 4; // for '] = '
50                }
51                if !self.check(v) {
52                    return false;
53                }
54            }
55            Value::Array(arr) => {
56                self.current_len += 1;
57                for v in arr {
58                    if !self.check(v) {
59                        return false;
60                    }
61                    self.current_len += 2;
62                }
63                self.current_len += 1;
64            }
65            Value::Null => {
66                self.current_len += 3; // "nil"
67            }
68        }
69        if self.current_len > self.target_len {
70            return false;
71        }
72        true
73    }
74}
75
76/// A dumper for Artemis AST scripts.
77pub struct Dumper<'a> {
78    current_indent: usize,
79    writer: Box<dyn Write + 'a>,
80    indent: Option<usize>,
81    max_line_width: usize,
82    current_line_width: usize,
83}
84
85impl<'a> Dumper<'a> {
86    /// Creates a new Dumper with the specified writer.
87    ///
88    /// default indent size is 4 spaces, and max line width is 100 characters.
89    pub fn new<W: Write + 'a>(writer: W) -> Self {
90        Dumper {
91            current_indent: 0,
92            writer: Box::new(writer),
93            indent: Some(4),
94            max_line_width: 100,
95            current_line_width: 0,
96        }
97    }
98
99    /// Sets the indent size for the dumper.
100    pub fn set_indent(&mut self, indent: usize) {
101        self.indent = Some(indent);
102    }
103
104    /// Disables indentation for the dumper.
105    pub fn set_no_indent(&mut self) {
106        self.indent = None;
107    }
108
109    /// Sets the maximum line width for the dumper.
110    pub fn set_max_line_width(&mut self, max_line_width: usize) {
111        self.max_line_width = max_line_width;
112    }
113
114    fn dump_f64(&mut self, f: &f64) -> std::io::Result<()> {
115        if f.fract() == 0.0 {
116            write!(self.writer, "{:.1}", f)
117        } else {
118            write!(self.writer, "{}", f)
119        }
120    }
121
122    /// Dumps the AST file to the writer.
123    pub fn dump(mut self, ast: &AstFile) -> std::io::Result<()> {
124        if self.indent.is_none() {
125            if let Some(astver) = ast.astver {
126                self.writer.write(b"astver=")?;
127                self.dump_f64(&astver)?;
128            }
129            if let Some(astname) = &ast.astname {
130                self.writer.write(b"\nastname = ")?;
131                if lua_str_contains_need_escape(astname) {
132                    self.writer.write(b"[[")?;
133                    self.writer.write(astname.as_bytes())?;
134                    self.writer.write(b"]]")?;
135                } else {
136                    self.writer.write(b"\"")?;
137                    self.writer.write(astname.as_bytes())?;
138                    self.writer.write(b"\"")?;
139                }
140            };
141            self.writer.write(b"\nast=")?;
142            self.dump_value(&ast.ast)?;
143        } else {
144            if let Some(astver) = ast.astver {
145                self.writer.write(b"astver = ")?;
146                self.dump_f64(&astver)?;
147            }
148            if let Some(astname) = &ast.astname {
149                self.writer.write(b"\nastname = ")?;
150                if lua_str_contains_need_escape(&astname) {
151                    self.writer.write(b"[[")?;
152                    self.writer.write(astname.as_bytes())?;
153                    self.writer.write(b"]]")?;
154                } else {
155                    self.writer.write(b"\"")?;
156                    self.writer.write(astname.as_bytes())?;
157                    self.writer.write(b"\"")?;
158                }
159            };
160            self.writer.write(b"\nast = ")?;
161            self.current_line_width = 6;
162            self.dump_value(&ast.ast)?;
163        }
164        self.writer.write(b"\n")?;
165        Ok(())
166    }
167
168    fn dump_value(&mut self, v: &Value) -> std::io::Result<()> {
169        if self.indent.is_none() {
170            match v {
171                Value::Float(f) => self.dump_f64(f)?,
172                Value::Int(i) => write!(self.writer, "{}", i)?,
173                Value::Str(s) => {
174                    if lua_str_contains_need_escape(s) {
175                        self.writer.write(b"[[")?;
176                        self.writer.write(s.as_bytes())?;
177                        self.writer.write(b"]]")?;
178                    } else {
179                        self.writer.write(b"\"")?;
180                        self.writer.write(s.as_bytes())?;
181                        self.writer.write(b"\"")?;
182                    }
183                }
184                Value::KeyVal((k, v)) => {
185                    if let Some(k) = k.as_str() {
186                        if lua_key_contains_need_escape(k) {
187                            self.writer.write(b"[\"")?;
188                            self.writer.write(k.as_bytes())?;
189                            self.writer.write(b"\"]=")?;
190                        } else {
191                            self.writer.write(k.as_bytes())?;
192                            self.writer.write(b"=")?;
193                        }
194                    } else {
195                        self.writer.write(b"[")?;
196                        self.dump_value(k)?;
197                        self.writer.write(b"]=")?;
198                    }
199                    self.dump_value(v)?;
200                }
201                Value::Array(arr) => {
202                    self.writer.write(b"{")?;
203                    for (i, v) in arr.iter().enumerate() {
204                        if i > 0 {
205                            self.writer.write(b",")?;
206                        }
207                        self.dump_value(v)?;
208                    }
209                    self.writer.write(b"}")?;
210                }
211                Value::Null => {
212                    self.writer.write(b"nil")?;
213                }
214            }
215        } else {
216            match v {
217                Value::Float(f) => self.dump_f64(f)?,
218                Value::Int(i) => write!(self.writer, "{}", i)?,
219                Value::Str(s) => {
220                    if lua_str_contains_need_escape(s) {
221                        self.writer.write(b"[[")?;
222                        self.writer.write(s.as_bytes())?;
223                        self.writer.write(b"]]")?;
224                    } else {
225                        self.writer.write(b"\"")?;
226                        self.writer.write(s.as_bytes())?;
227                        self.writer.write(b"\"")?;
228                    }
229                }
230                Value::KeyVal((k, v)) => {
231                    if let Some(k) = k.as_str() {
232                        let bytes = k.as_bytes();
233                        if lua_key_contains_need_escape(k) {
234                            self.writer.write(b"[\"")?;
235                            self.writer.write(bytes)?;
236                            self.writer.write(b"\"] = ")?;
237                            self.current_line_width += bytes.len() + 7;
238                        } else {
239                            self.writer.write(bytes)?;
240                            self.writer.write(b" = ")?;
241                            self.current_line_width += bytes.len() + 3;
242                        }
243                    } else {
244                        self.writer.write(b"[")?;
245                        self.current_line_width += 1;
246                        self.dump_value(k)?;
247                        self.writer.write(b"] = ")?;
248                        self.current_line_width += 4; // for '] = '
249                    };
250                    if v.is_array() {
251                        let tlen = self.current_line_width + self.current_indent;
252                        if tlen < self.max_line_width {
253                            let mut checker = LenChecker::new(self.max_line_width - tlen);
254                            if checker.check(v) {
255                                self.dump_value_in_one(v)?;
256                                return Ok(());
257                            }
258                        }
259                    }
260                    self.dump_value(v)?;
261                }
262                Value::Array(a) => {
263                    let tlen = self.current_line_width + self.current_indent;
264                    if tlen < self.max_line_width {
265                        let mut checker = LenChecker::new(self.max_line_width - tlen);
266                        if checker.check(v) {
267                            self.dump_value_in_one(v)?;
268                            return Ok(());
269                        }
270                    }
271                    self.writer.write(b"{\n")?;
272                    self.current_indent += self.indent.unwrap();
273                    for (i, v) in a.iter().enumerate() {
274                        if i > 0 {
275                            self.writer.write(b",\n")?;
276                        }
277                        self.dump_indent()?;
278                        self.current_line_width = 0;
279                        self.dump_value(v)?;
280                    }
281                    self.current_indent -= self.indent.unwrap();
282                    self.writer.write(b",\n")?;
283                    self.dump_indent()?;
284                    self.writer.write(b"}")?;
285                }
286                Value::Null => {
287                    self.writer.write(b"nil")?;
288                }
289            }
290        }
291        Ok(())
292    }
293
294    fn dump_indent(&mut self) -> std::io::Result<()> {
295        for _ in 0..self.current_indent {
296            self.writer.write(b" ")?;
297        }
298        Ok(())
299    }
300
301    fn dump_value_in_one(&mut self, v: &Value) -> std::io::Result<()> {
302        match v {
303            Value::Float(f) => self.dump_f64(f)?,
304            Value::Int(i) => write!(self.writer, "{}", i)?,
305            Value::Str(s) => {
306                if lua_str_contains_need_escape(s) {
307                    self.writer.write(b"[[")?;
308                    self.writer.write(s.as_bytes())?;
309                    self.writer.write(b"]]")?;
310                } else {
311                    self.writer.write(b"\"")?;
312                    self.writer.write(s.as_bytes())?;
313                    self.writer.write(b"\"")?;
314                }
315            }
316            Value::KeyVal((k, v)) => {
317                if let Some(k) = k.as_str() {
318                    if lua_key_contains_need_escape(k) {
319                        self.writer.write(b"[\"")?;
320                        self.writer.write(k.as_bytes())?;
321                        self.writer.write(b"\"]=")?;
322                    } else {
323                        self.writer.write(k.as_bytes())?;
324                        self.writer.write(b"=")?;
325                    }
326                } else {
327                    self.writer.write(b"[")?;
328                    self.dump_value_in_one(k)?;
329                    self.writer.write(b"]=")?;
330                };
331                self.dump_value_in_one(v)?;
332            }
333            Value::Array(arr) => {
334                self.writer.write(b"{")?;
335                for (i, v) in arr.iter().enumerate() {
336                    if i > 0 {
337                        self.writer.write(b", ")?;
338                    }
339                    self.dump_value_in_one(v)?;
340                }
341                self.writer.write(b"}")?;
342            }
343            Value::Null => {
344                self.writer.write(b"nil")?;
345            }
346        }
347        Ok(())
348    }
349}